// Copyright (c) 2010 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// NOTE: this file is translated from reboot_mode.py
//
// A script to manipulate ChromeOS CMOS reboot Field.
//
// A few bits in a byte in CMOS (called RDB, short for 'Reboot Data Byte'
// hereafter) are dedicated to information exchange between Linux and BIOS.
//
// The CMOS contents are available through /dev/nvrom. The location of RDB is
// reported by BIOS through ACPI. The few bits in RDB operated on by this script
// are defined in the all_mode_fields dictionary below along with their
// functions.
//
// When invoked without parameters this script prints values of the fields. When
// invoked with any of the bit field names as parameter(s), the script sets the
// bit(s) as required.
//
// Additional parameters allow to specify alternative ACPI and NVRAM files for
// testing.

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <unistd.h>
#include <sys/param.h>
#include <stdint.h>
#include <errno.h>
#include <stdarg.h>
#include <getopt.h>

#include <string>
#include <algorithm>
#include <vector>
#include <map>

#define DEFAULT_ACPI_FILE       "/sys/devices/platform/chromeos_acpi/CHNV"
#define DEFAULT_NVRAM_FILE      "/dev/nvram"

///////////////////////////////////////////////////////////////////////
// Utilities for quick python-C translation

using std::string;

// Works like throwing an exception - directly exit here.
static void RebootModeError(const char *message, ...) {
  va_list args;
  va_start(args, message);
  vfprintf(stderr, message, args);
  fprintf(stderr, "\n");
  va_end(args);
  exit(1);
}

// return a string which is upper case of input parameter.
static string str_to_upper(const string str) {
  string newstr = str;
  std::transform(str.begin(), str.end(), newstr.begin(), toupper);
  return newstr;
}

///////////////////////////////////////////////////////////////////////

typedef std::map<string, uint8_t> RDB_BitField;

struct RebootModeConfig {
  // file path
  string acpi_file;
  string nvram_file;

  // Offset of RDB in NVRAM
  int rdb_offset;

  // NVRAM contents read on startup
  uint8_t rdb_value;

  // A dictionary of fields to be updated as requested by the command line
  // parameters
  RDB_BitField fields_to_set;

  // All bitfields in RDB this script provides access to. Field names prepended
  // by `--' become this script's command line options
  RDB_BitField all_mode_fields;

  // write constructure here to ease python translation
  RebootModeConfig() {
    acpi_file = DEFAULT_ACPI_FILE;
    nvram_file = DEFAULT_NVRAM_FILE;
    rdb_offset = 0;
    all_mode_fields["recovery"] = 0x80;
    all_mode_fields["debug_reset"] = 0x40;
    all_mode_fields["try_firmware_b"] = 0x20;
  }
} cfg;  // this is a special global object to ease python translation.

// Process bit field name command line parameter.
//
// Verifies that the value is in range (0 or 1) and adds the appropriate
// element to fields_to_set dictionary. Should the same field specified in the
// command line more than once, only the last value will be used.
//
// Raises:
//   RebootModeError in case the parameter value is out of range.
void OptionHandler(const char *opt_str, const char *value_str) {
  const char *key = opt_str;
  int value = atoi(value_str);

  if (!value && *value_str != '0')
    RebootModeError("--%s needs numeric argument value", key);

  if (value < 0 || value > 1)
    RebootModeError("--%s should be either 0 or 1", key);

  if (cfg.all_mode_fields.find(key) == cfg.all_mode_fields.end())
    RebootModeError("INTERNAL ERROR: UNKNOWN FIELD: %s", key);

  cfg.fields_to_set[key] = value;
}

void GetRDBOffset() {
  const char *acpi_file = cfg.acpi_file.c_str();

  FILE *f = fopen(acpi_file, "rb");
  if (!f) {
    perror(acpi_file);
    RebootModeError("Cannot read acpi_file: %s", acpi_file);
  }

  char buffer[PATH_MAX];  // this should be large enough
  if (fgets(buffer, sizeof(buffer), f) == NULL) {
    perror(acpi_file);
    RebootModeError("Trying to read %s but failed.", acpi_file);
  }

  cfg.rdb_offset = atoi(buffer);
  if (!cfg.rdb_offset && *buffer != '0') {  // consider this as error
    RebootModeError("%s contents are corrupted", acpi_file);
  }
}

void ReadNvram() {
  FILE *f = fopen(cfg.nvram_file.c_str(), "rb");
  if (!f) {
    if (errno == 2)
      RebootModeError("%s does not exist. Is nvram module installed?",
                      cfg.nvram_file.c_str());
    if (errno == 13)
      RebootModeError("Access to %s denied. Are you root?",
                      cfg.nvram_file.c_str());
    // all other error
    perror(cfg.nvram_file.c_str());
  }

  if (fseek(f, cfg.rdb_offset, SEEK_SET) == -1 ||
      fread(&cfg.rdb_value, 1, 1, f) != 1) {
    perror(cfg.nvram_file.c_str());
    RebootModeError("Failed reading mode byte from %s", cfg.nvram_file.c_str());
  }
  fclose(f);
}

void PrintCurrentMode() {
  printf("Current reboot mode settings:\n");
  for (RDB_BitField::iterator i = cfg.all_mode_fields.begin();
       i != cfg.all_mode_fields.end();
       ++i) {
    printf("%-15s: %d\n", i->first.c_str(),
           (cfg.rdb_value & i->second) ? 1 : 0);
  }
}

void SetNewMode() {
  uint8_t new_mode = 0;
  uint8_t updated_bits_mask = 0;

  for (RDB_BitField::iterator i = cfg.fields_to_set.begin();
       i != cfg.fields_to_set.end();
       ++i) {
    string opt = i->first;
    uint8_t value = i->second;
    assert(cfg.all_mode_fields.find(opt) != cfg.all_mode_fields.end());

    uint8_t mask = cfg.all_mode_fields[opt];
    if (value)
      new_mode |= mask;

    updated_bits_mask |= mask;
  }

  if ((cfg.rdb_value & updated_bits_mask) == new_mode) {
    printf("No update required\n");
    return;
  }

  FILE *f = fopen(cfg.nvram_file.c_str(), "rb+");
  fseek(f, cfg.rdb_offset, SEEK_SET);
  new_mode |= (cfg.rdb_value & ~updated_bits_mask);
  if (fwrite(&new_mode, 1, 1, f) != 1) {
    RebootModeError("Failed writing mode byte to %s", cfg.nvram_file.c_str());
  }
  fclose(f);
}

const char *__name__;

void usage_help_exit(int err) {
  printf("Usage: %s [options]\n"
         "\n"
         "Options:\n"
         "  -h, --help\t\tshow this help message and exit\n"
         "  --acpi_file=ACPI_FILE\n"
         "  --nvram_file=NVRAM_FILE\n"
         , __name__);

  // list all possible fields
  for (RDB_BitField::iterator i = cfg.all_mode_fields.begin();
       i != cfg.all_mode_fields.end();
       ++i) {
    printf("  --%s=%s\n",
           i->first.c_str(),
           str_to_upper(i->first).c_str());
  }
  exit(err);
}

int main(int argc, char *argv[]) {
  __name__ = argv[0];

  struct option optv = {0};
  std::vector<struct option> longopts;
  int max_field_index = 0;

  optv.has_arg = 0;
  optv.val = 0;
  optv.name = "help";           // value = 0
  longopts.push_back(optv);
  optv.has_arg = 1;

  optv.val++;
  optv.name = "acpi_file";      // value = 1
  longopts.push_back(optv);
  optv.val++;
  optv.name = "nvram_file";     // value = 2
  longopts.push_back(optv);

  // add all known RDB fields
  for (RDB_BitField::iterator i = cfg.all_mode_fields.begin();
       i != cfg.all_mode_fields.end();
       ++i) {
    optv.name = i->first.c_str();
    optv.val++;
    longopts.push_back(optv);
  }
  max_field_index = optv.val;

  // add a zero for closing options
  optv.name = NULL;
  optv.val = optv.has_arg = 0;
  longopts.push_back(optv);

  int optc;
  while ((optc = getopt_long(argc, argv, "h", &longopts.front(), NULL)) != -1) {
    switch (optc) {
      case 1:
        cfg.acpi_file = optarg;
        break;

      case 2:
        cfg.nvram_file = optarg;
        break;

      default:
        // RDB fields?
        // printf("optc: %d %s %s\n", optc, longopts[optc].name, optarg);
        if (optc > 0 && optc <= max_field_index)
          OptionHandler(longopts[optc].name, optarg);
        else
          usage_help_exit(1);
        break;
    }
  }

  // currently no other non-dashed arguments allowed.
  if (optind != argc)
    usage_help_exit(1);

  GetRDBOffset();
  ReadNvram();

  if (!cfg.fields_to_set.empty()) {
    SetNewMode();
  } else {
    PrintCurrentMode();
  }
  return 0;
}
